package com.yanghui.dtracing.core.opentracing;

import io.opentracing.Scope;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

public class ScopeContext {

    private static ThreadLocal<Map<String,Scope>> SCOPE_THREAD_LOCAL = new ThreadLocal<Map<String,Scope>>(){
        @Override
        protected Map<String, Scope> initialValue() {
            return new HashMap<>();
        }
    };

    private static ThreadLocal<AtomicInteger> AUTO_LOCAL = new ThreadLocal<>();

    public static void putIncrement(AtomicInteger increment){
        AUTO_LOCAL.set(increment);
    }

    public static AtomicInteger getIncrement(){
        return AUTO_LOCAL.get();
    }

    public static int increment(int initial){
        AtomicInteger increment = AUTO_LOCAL.get();
        if(increment == null){
            increment = new AtomicInteger(initial);
            AUTO_LOCAL.set(increment);
        }
        return increment.incrementAndGet();
    }

    public static void put(String spanId,Scope scope){
        SCOPE_THREAD_LOCAL.get().put(spanId,scope);
    }

    public static void close(String spanId){
        Map<String,Scope> scopeMap = SCOPE_THREAD_LOCAL.get();
        if(null != scopeMap && scopeMap.get(spanId) != null){
            scopeMap.get(spanId).close();
            scopeMap.remove(spanId);
        }
        if(Tracing.getTracer().activeSpan() == null){
            AUTO_LOCAL.remove();
        }
    }
}
